# Makefile for building dev/MACA kernels
# Collects all the make commands in one file but each file also
# has the compile and run commands in the header comments section.
# Location of the MACA Toolkit
MACA_PATH       ?= /opt/maca

# Compiler
MXCC = $(MACA_PATH)/mxgpu_llvm/bin/mxcc
#MXCC := mxcc

# Internal flags
MXCC_FLAGS := -x maca -lmcruntime
MXCC_FLAGS_EXT = -lmcblas -lmcblasLt -std=c++17
MPI_PATHS = -I/usr/lib/x86_64-linux-gnu/openmpi/include -L/usr/lib/x86_64-linux-gnu/openmpi/lib/

# Location of head files
MACA_INCLUDE ?= $(MACA_PATH)include
MY_INCLUDE_PATH ?= ../../include

# Default rule for our CUDA files
%: %.maca
	$(MXCC) $(MXCC_FLAGS) $(MXCC_FLAGS_EXT) $< -o $@ -I$(MACA_INCLUDE) --maca-path=$(MACA_PATH)

##############################################################################################
# Target rules
TARGETS = adamw attention_backward attention_forward classifier_fused crossentropy_forward crossentropy_softmax_backward encoder_backward encoder_forward gelu_backward gelu_forward layernorm_backward layernorm_forward matmul_backward matmul_backward_bias matmul_forward mccl_all_reduce residual_forward softmax_forward trimat_forward fused_residual_forward  global_norm permute

all: $(TARGETS)

# Individual targets: forward pass
attention_forward: attention_forward.maca
classifier_fused: classifier_fused.maca
crossentropy_forward: crossentropy_forward.maca
encoder_forward: encoder_forward.maca
gelu_forward: gelu_forward.maca
layernorm_forward: layernorm_forward.maca
fused_residual_forward: fused_residual_forward.maca
residual_forward: residual_forward.maca
softmax_forward: softmax_forward.maca
trimat_forward: trimat_forward.maca
# matmul fwd/bwd also uses OpenMP (optionally) and cuBLASLt libs
matmul_forward: matmul_forward.maca
	$(MXCC) $(MXCC_FLAGS) $(MXCC_FLAGS_EXT) -Xcompiler -fopenmp matmul_forward.maca -o matmul_forward

# Individual targets: backward pass
attention_backward: attention_backward.maca
crossentropy_softmax_backward: crossentropy_softmax_backward.maca
encoder_backward: encoder_backward.maca
gelu_backward: gelu_backward.maca
layernorm_backward: layernorm_backward.maca
matmul_backward_bias: matmul_backward_bias.maca
matmul_backward: matmul_backward.maca
	$(MXCC) $(MXCC_FLAGS) $(MXCC_FLAGS_EXT) -Xcompiler -fopenmp matmul_backward.maca -o matmul_backward

# Update kernels
adamw: adamw.maca
global_norm: global_norm.maca

permute: permute.maca

# MCCL communication kernels
mccl_all_reduce: mccl_all_reduce.maca
	$(MXCC) -lmpi -lmccl $(MXCC_FLAGS_EXT) $(MPI_PATHS) mccl_all_reduce.maca -o mccl_all_reduce
# Run all targets
run_all: all
	@for target in $(TARGETS); do \
		echo "\n========================================"; \
		echo "Running $$target ..."; \
		echo "========================================\n"; \
		./$$target; \
	done

# Clean up
clean:
	rm -f $(TARGETS)